import numpy as np
import os
from PIL import Image, ImageDraw
from torch.utils.data import DataLoader
from transformers import DetrImageProcessor, DetrForObjectDetection, AutoModelForObjectDetection
from matplotlib import pyplot as plt
from dataset_create import CocoDetection

processor = DetrImageProcessor.from_pretrained("/data/models/detr-resnet-50-train")
train_dataset = CocoDetection(img_folder='/data/datasets/coco/train2017', processor=processor)
val_dataset = CocoDetection(img_folder='/data/datasets/coco/val2017', processor=processor, train=False)

print("Number of training examples:", len(train_dataset))
print("Number of validation examples:", len(val_dataset))

def showData():
    image_ids = train_dataset.coco.getImgIds()
    # let's pick a random image
    image_id = image_ids[np.random.randint(0, len(image_ids))]
    print('Image n°{}'.format(image_id))
    image = train_dataset.coco.loadImgs(image_id)[0]
    image = Image.open(os.path.join('/data/datasets/coco/train2017', image['file_name']))
    annotations = train_dataset.coco.imgToAnns[image_id]
    draw = ImageDraw.Draw(image, "RGBA")

    cats = train_dataset.coco.cats
    id2label = {k: v['name'] for k, v in cats.items()}

    for annotation in annotations:
        box = annotation['bbox']
        class_idx = annotation['category_id']
        x, y, w, h = tuple(box)
        draw.rectangle((x, y, x + w, y + h), outline='red', width=1)
        draw.text((x, y), id2label[class_idx], fill='white')

    image.show()
    plt.show()

if __name__ == '__main__':
    showData()
